package com.example.rbacmanage.config;

import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import org.apache.commons.collections4.CollectionUtils;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.http.server.ServletServerHttpResponse;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.util.List;
import java.util.Map;
/**
 * <h2></h2>
 * <p>
 * websocket接口地址要在开发权限
 * </p>
 *
 * @createTime 2022年07月19日 9:57?上午
 */
@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(webSocketHandler(), "/websocket")
                .addInterceptors(new HandshakeInterceptor() {
                    @Override
                    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
                        List<String> tokens = request.getHeaders().get("Sec-WebSocket-Protocol");
                        if (CollectionUtils.isNotEmpty(tokens)) {
                            String token = tokens.get(0);
                        }
                        return true;
                    }

                    @Override
                    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {
                        // 获取 Servlet 的 HttpServletRequest 和 HttpServletResponse 对象
                        // httpRequest 可以获取 HTTP协议升级前 请求报文的信息，如 header中的键值对等
                        // httpResponse 可以设置 HTTP响应 的相关信息，如状态码、ContentType、header信息等
                        HttpServletRequest httpRequest = ((ServletServerHttpRequest) request).getServletRequest();
                        HttpServletResponse httpResponse = ((ServletServerHttpResponse) response).getServletResponse();
                        if (httpRequest.getHeader("Sec-WebSocket-Protocol") != null) {
                            httpResponse.addHeader("Sec-WebSocket-Protocol", httpRequest.getHeader("Sec-WebSocket-Protocol"));
                        }
                    }
                })
                .setAllowedOrigins("*"); // 允许跨域访问
    }

    @Bean
    public WebSocketHandler webSocketHandler() {
        // 使用自定义的WebSocket处理器
        return new WebSocketCustomHandler();
    }

}
